from anomaly.datasets import MVTecDataset_knn, mvtec_classes
from anomaly.models import PatchCoreK
# from anomaly.utils.knn_utils import write_results

# seeds
import torch
import random
import numpy as np

from argparse import ArgumentParser

torch.manual_seed(0)
random.seed(0)
np.random.seed(0)

import warnings # for some torch warnings regarding depreciation
warnings.filterwarnings("ignore")

ALL_CLASSES = mvtec_classes()
ALLOWED_METHODS = ["patchcore"]   # same name padim 


def parser_args():
    parser = ArgumentParser()
    parser.add_argument('--method', default="patchcore")
    parser.add_argument("--dataset", default="all", help="Dataset, defaults to all datasets.")
    parser.add_argument("--weights_dir", default="/home/ops/lnq/anomaly_lab/knn_patchcore/weights", help="model dir to save")
    parser.add_argument("--result_dir", default="/home/ops/lnq/anomaly_lab/knn_patchcore/result", help="results of heatmap region and scores to save ")
    args = parser.parse_args() 
    
    return args


def run_model(args, method, classes):
    # 所有类，不同正常
    results = {}
    for cls  in classes:
        method == "patchcore"
        model = PatchCoreK(args, f_coreset=0.01, backbone="wide_resnet50_2")

        print(f"Running {method} on {cls} dataset.")
        train_ds, test_ds = MVTecDataset_knn(cls).load()

        print("Training ...")
        model.fit(train_ds)   
        print("Testing ...")
        image_rocauc, pixel_rocauc = model.evaluate(test_ds)   
        
        print(f"Test results {cls} - image_rocauc: {image_rocauc:.2f}, pixel_rocauc: {pixel_rocauc:.2f}")
        results[cls] = [float(image_rocauc), float(pixel_rocauc)]

    image_results = [v[0] for _, v in results.items()]
    average_image_roc_auc = sum(image_results)/len(image_results)
    image_results = [v[1] for _, v in results.items()]
    average_pixel_roc_auc = sum(image_results)/len(image_results)


    total_results = {
        "per_class_results": results,
        "average image rocauc": average_image_roc_auc,
        "average pixel rocauc": average_pixel_roc_auc,
        "model parameters": model.get_parameters(),
    }
    

    return total_results


def knns_train(args): 
    method = args.method
    dataset = args.dataset
    if dataset == "all":
        dataset = ALL_CLASSES
    else:
        assert dataset in ALL_CLASSES, "Dataset does not exist."
        dataset = [dataset]

    method = method.lower()
    assert method in ALLOWED_METHODS, f"Select from {ALLOWED_METHODS}."

    total_results = run_model(args, method, dataset)

    print('knn_patchcore: ', total_results)

    # write_results(total_results, method)
    

if __name__ == "__main__":

    args = parser_args()
    knns_train(args)